import numpy as np
import jax
import equinox as eqx
import jax.numpy as jnp
from jaxtyping import Array, Float, Int, Bool, PyTree
from models.modules.utils import process_t

def hawkes_intensities_points(hawkes, ts: Float[Array, "T"], marks: Int[Array, "T"], t_max: float, num_samples: int):
    xs = list()
    ys = list()
    T = len(ts)
    ts_ = jnp.concatenate((ts, jnp.asarray([t_max, ])), -1) # (T+1, )
    dts_ = ts_[1:] - ts_[:-1] # (T, )
    for i in range(T):
        dt_cur = dts_[i]
        samples = np.linspace(0., dt_cur, num_samples) # (num_samples, )
        mask = jnp.asarray([True, ]*(i+1)+[False, ]*(T-i-1), dtype=bool)
        intensities = jax.vmap(eqx.filter_jit(hawkes.intensities_at), (None, None, None, 0))(ts, marks, mask, samples) # (num_samples, num_types)
        xs.append(ts[i]+samples)
        ys.append(intensities)
    xs = np.concatenate(xs)
    ys = np.concatenate(ys, 0)
    return xs, ys # xs: (total_sample_num, ) ys: (total_sample_num, num_types)

def poisson_intensities_points(poisson, ts, t_max):
    xs = np.asarray([ts[0], t_max])
    ys = np.stack((poisson.params, poisson.params), 0) # (2, num_types)
    return xs, ys

def model_intensities_points(model, ts: Float[Array, "T"], marks: Int[Array, "T"], t_scale, t_max, num_samples, key):
    model = eqx.nn.inference_mode(model)
    xs = list()
    ys = list()
    key1, key2 = jax.random.split(key)

    ts_ = process_t(ts, t_scale)
    states = model.encode(ts, marks, jnp.ones_like(ts, dtype=bool), key1) # (T, ...)
    ts_ = jnp.concatenate((ts_, jnp.asarray([(t_max-ts[0])/t_scale, ])), -1) # (T+1, )
    dts_ = ts_[1:] - ts_[:-1] # (T, )
    T = dts_.shape[0]
    for i in range(T):
        samples = jnp.linspace(0., dts_[i], num_samples) # (num_samples, )
        cur_key, key2 = jax.random.split(key2)
        state_cur = jax.tree.map(lambda x: x[i], states)
        intensities = model.intensities_at(state_cur, samples, cur_key) # (num_samples, num_types)
        intensities = intensities / t_scale # (num_samples, num_types)
        x = ts[i] + samples * t_scale # (num_samples, )
        xs.append(x)
        ys.append(intensities)
    xs = np.concatenate(xs)
    ys = np.concatenate(ys, 0)
    return xs, ys # xs: (total_sample_num, ) ys: (total_sample_num, num_types)

def event_scatter_points(ts, marks):
    xs = list()
    num_types = marks.max() + 1
    for i in range(num_types):
        mask = marks==i
        x = np.asarray(ts[mask])
        xs.append(x)
    return xs


